# -*- coding: UTF-8 -*-
# @File : embeddings_builder.py
# @IDE     : VScode
# @Author  : zhonggc
# @Date    : 2024/05/28 15:28:07

import os
from zhipuai_embeddings import ZhipuAIEmbeddings
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_community.document_loaders import UnstructuredMarkdownLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma


class EmbeddingsBuilder:
    loader_utils = {
        "pdf": PyMuPDFLoader,
        "md": UnstructuredMarkdownLoader,
    }

    def __init__(self, doc_folder_path: str) -> None:
        """初始化

        Arguments:
            doc_folder_path -- 指定知识库文件夹路径
        """
        self.doc_folder_path = doc_folder_path
        self.file_path = []
        self.loaders = []
        self.texts = []

    def get_file_paths(self) -> None:
        """获取文件路径"""
        for root, dirs, files in os.walk(self.doc_folder_path):
            for file in files:
                self.file_path.append(os.path.join(root, file))

    def load_files(self) -> None:
        """加载文件"""
        for file in self.file_path:
            # 根据文件类型选择加载器
            file_type = file.split(".")[-1]
            if file_type in self.loader_utils:
                self.loaders.append(self.loader_utils[file_type](file))

    def save_text(self) -> None:
        """读取文件内容保存到text"""
        for loader in self.loaders:
            self.texts.extend(loader.load())

    def text_split(self) -> None:
        """文本分割"""
        text_spliter = RecursiveCharacterTextSplitter(
            chunk_size=1000, chunk_overlap=100
        )
        return text_spliter.split_documents(self.texts)

    def build_chroma(self, documents):
        """构建chroma向量库"""
        embedding = ZhipuAIEmbeddings()
        # 定义持久化路径
        persist_directory = "./src/ChatBot/DataBase/chroma"
        if len(documents) == 0:
            print("No documents to embed.")
            return
        self.vectordb = Chroma.from_documents(
            documents=documents,
            embedding=embedding,
            persist_directory=persist_directory,
        )
        self.vectordb.persist()

    def build_embeddings(self):
        """构建embeddings"""
        self.get_file_paths()
        self.load_files()
        self.save_text()
        documents = self.text_split()
        self.build_chroma(documents)
        print("Embeddings build success.")


def main():
    file_path = r"E:\python\Docs"
    embeddings_builder = EmbeddingsBuilder(file_path)
    embeddings_builder.build_embeddings()


if __name__ == "__main__":
    main()
